"""
数据库连接管理器
支持MySQL、PostgreSQL等多种数据库的连接和查询
"""
import time
import logging
import sqlparse
from typing import Dict, List, Tuple, Any, Optional
from contextlib import contextmanager
from django.conf import settings
from sqlalchemy import create_engine, text, MetaData, inspect
from sqlalchemy.exc import SQLAlchemyError, OperationalError, ProgrammingError
from sqlalchemy.pool import QueuePool
import pymysql
import psycopg2

logger = logging.getLogger(__name__)


class DatabaseConnectionError(Exception):
    """数据库连接异常"""
    pass


class SQLExecutionError(Exception):
    """SQL执行异常"""
    pass


class DatabaseManager:
    """数据库连接管理器"""
    
    # 连接池缓存
    _engines = {}
    
    # 危险SQL关键字
    DANGEROUS_KEYWORDS = [
        'DROP', 'DELETE', 'TRUNCATE', 'ALTER', 'CREATE', 'INSERT', 'UPDATE',
        'GRANT', 'REVOKE', 'FLUSH', 'RESET', 'SHUTDOWN', 'KILL'
    ]
    
    # 只读SQL关键字
    READONLY_KEYWORDS = ['SELECT', 'SHOW', 'DESCRIBE', 'DESC', 'EXPLAIN']
    
    @classmethod
    def get_connection_string(cls, db_instance) -> str:
        """根据数据库实例生成连接字符串"""
        db_type = db_instance.database_type.lower()
        host = db_instance.host
        port = db_instance.port
        username = db_instance.username
        password = db_instance.password
        database = db_instance.database_name
        
        if db_type == 'mysql':
            # MySQL连接字符串处理，如果数据库名称为空，则不指定数据库
            if database:
                return f"mysql+pymysql://{username}:{password}@{host}:{port}/{database}?charset=utf8mb4"
            else:
                return f"mysql+pymysql://{username}:{password}@{host}:{port}/?charset=utf8mb4"
        elif db_type == 'postgresql':
            # PostgreSQL连接字符串处理，如果数据库名称为空，则连接到默认数据库
            if database:
                return f"postgresql+psycopg2://{username}:{password}@{host}:{port}/{database}"
            else:
                return f"postgresql+psycopg2://{username}:{password}@{host}:{port}/postgres"
        elif db_type == 'sqlite':
            # SQLite数据库路径处理
            if database == ':memory:':
                return "sqlite:///:memory:"
            else:
                return f"sqlite:///{database}"
        else:
            raise DatabaseConnectionError(f"不支持的数据库类型: {db_type}")
    
    @classmethod
    def get_engine(cls, db_instance):
        """获取数据库引擎（带连接池）"""
        # 使用实例ID、更新时间和数据库名作为缓存键，确保配置变更时重新创建引擎
        # 包含数据库名是为了支持同一实例连接不同数据库的情况
        db_name = db_instance.database_name or 'default'
        cache_key = f"{db_instance.id}_{db_instance.updated_at.timestamp()}_{db_name}"
        
        if cache_key not in cls._engines:
            try:
                connection_string = cls.get_connection_string(db_instance)
                
                # 创建引擎，配置连接池
                if db_instance.database_type.lower() == 'sqlite':
                    # SQLite使用单线程连接
                    engine = create_engine(
                        connection_string,
                        echo=False,
                        connect_args={'check_same_thread': False}
                    )
                else:
                    # 其他数据库使用连接池
                    engine = create_engine(
                        connection_string,
                        poolclass=QueuePool,
                        pool_size=5,
                        max_overflow=10,
                        pool_timeout=30,
                        pool_recycle=3600,
                        echo=False  # 生产环境设为False
                    )
                
                # 测试连接
                with engine.connect() as conn:
                    conn.execute(text("SELECT 1"))
                
                cls._engines[cache_key] = engine
                logger.info(f"数据库连接创建成功: {db_instance.name}")
                
            except Exception as e:
                logger.error(f"数据库连接失败: {db_instance.name}, 错误: {str(e)}")
                raise DatabaseConnectionError(f"连接数据库失败: {str(e)}")
        
        return cls._engines[cache_key]
    
    @classmethod
    def test_connection(cls, db_instance) -> Dict[str, Any]:
        """测试数据库连接"""
        try:
            start_time = time.time()
            engine = cls.get_engine(db_instance)
            
            with engine.connect() as conn:
                result = conn.execute(text("SELECT 1 as test"))
                result.fetchone()
            
            connection_time = time.time() - start_time
            
            return {
                'success': True,
                'message': '连接成功',
                'connection_time': round(connection_time, 3),
                'database_type': db_instance.database_type,
                'version': cls._get_database_version(engine, db_instance.database_type)
            }
            
        except Exception as e:
            logger.error(f"数据库连接测试失败: {db_instance.name}, 错误: {str(e)}")
            return {
                'success': False,
                'message': f'连接失败: {str(e)}',
                'connection_time': None,
                'database_type': db_instance.database_type,
                'version': None
            }
    
    @classmethod
    def _get_database_version(cls, engine, db_type: str) -> Optional[str]:
        """获取数据库版本"""
        try:
            with engine.connect() as conn:
                if db_type.lower() == 'mysql':
                    result = conn.execute(text("SELECT VERSION() as version"))
                elif db_type.lower() == 'postgresql':
                    result = conn.execute(text("SELECT version() as version"))
                else:
                    return None
                
                row = result.fetchone()
                return row[0] if row else None
        except:
            return None
    
    @classmethod
    def parse_sql(cls, sql_content: str) -> List[str]:
        """解析SQL语句，分割多个语句"""
        try:
            # 使用sqlparse解析SQL
            statements = sqlparse.split(sql_content)
            # 过滤空语句
            return [stmt.strip() for stmt in statements if stmt.strip()]
        except Exception as e:
            logger.error(f"SQL解析失败: {str(e)}")
            return [sql_content.strip()]
    
    @classmethod
    def analyze_sql_type(cls, sql: str) -> Dict[str, Any]:
        """分析SQL类型和安全性"""
        sql_upper = sql.upper().strip()

        # 解析SQL语句
        try:
            parsed = sqlparse.parse(sql)[0]
            first_token = None

            # 查找第一个关键字token
            for token in parsed.flatten():
                if token.ttype is sqlparse.tokens.Keyword and token.value.upper() in (cls.READONLY_KEYWORDS + cls.DANGEROUS_KEYWORDS):
                    first_token = token.value.upper()
                    break

            # 如果没找到，尝试简单分割
            if not first_token:
                words = sql_upper.split()
                if words:
                    first_token = words[0]

        except:
            words = sql_upper.split()
            first_token = words[0] if words else ''

        # 判断SQL类型
        is_readonly = first_token in cls.READONLY_KEYWORDS
        is_dangerous = first_token in cls.DANGEROUS_KEYWORDS

        # 特殊检查
        is_dml = first_token in ['INSERT', 'UPDATE', 'DELETE']
        is_ddl = first_token in ['CREATE', 'ALTER', 'DROP', 'TRUNCATE']
        is_dcl = first_token in ['GRANT', 'REVOKE']

        return {
            'sql_type': first_token,
            'is_readonly': is_readonly,
            'is_dangerous': is_dangerous,
            'is_dml': is_dml,
            'is_ddl': is_ddl,
            'is_dcl': is_dcl,
            'requires_confirmation': is_dangerous and first_token in ['DROP', 'TRUNCATE', 'DELETE']
        }
    
    @classmethod
    def execute_query(cls, db_instance, sql_content: str, user, limit: int = 1000) -> Dict[str, Any]:
        """执行SQL查询"""
        start_time = time.time()
        
        try:
            # 解析SQL语句
            statements = cls.parse_sql(sql_content)
            if not statements:
                raise SQLExecutionError("SQL语句为空")

            # 处理多SQL语句执行
            if len(statements) == 1:
                # 单条语句，使用原有逻辑
                return cls._execute_single_query(db_instance, statements[0], user, limit, start_time)
            else:
                # 多条语句，逐条执行
                return cls._execute_multiple_queries(db_instance, statements, user, limit, start_time)

        except (OperationalError, ProgrammingError) as e:
            logger.error(f"SQL执行错误: {str(e)}")
            raise SQLExecutionError(f"SQL执行错误: {str(e)}")
        except SQLAlchemyError as e:
            logger.error(f"数据库操作错误: {str(e)}")
            raise SQLExecutionError(f"数据库操作错误: {str(e)}")
        except Exception as e:
            logger.error(f"未知错误: {str(e)}")
            raise SQLExecutionError(f"执行失败: {str(e)}")

    @classmethod
    def _execute_single_query(cls, db_instance, sql_content: str, user, limit: int, start_time: float) -> Dict[str, Any]:
        """执行单条SQL查询"""
        try:
            # 检查是否为USE语句+查询语句的组合
            use_database = None
            db_instance_to_use = db_instance
            sql = sql_content.strip()

            # 检查是否是USE语句
            sql_upper = sql.upper()
            if sql_upper.startswith('USE '):
                # 从原始语句中提取数据库名（保持原始大小写）
                use_match = sql[4:].strip()  # 移除'USE '
                # 移除分号、反引号和引号
                use_database = use_match.rstrip(';').strip('`').strip("'").strip('"')
                logger.info(f"检测到USE语句，切换到数据库: {use_database}")

                # 创建一个临时的数据库实例副本，修改数据库名
                from copy import copy
                db_instance_to_use = copy(db_instance)
                db_instance_to_use.database_name = use_database

                # USE语句执行成功，返回结果
                execution_time = time.time() - start_time
                return {
                    'success': True,
                    'columns': [],
                    'rows': [],
                    'row_count': 0,
                    'affected_rows': 0,
                    'execution_time': round(execution_time, 3),
                    'sql_type': 'USE',
                    'is_readonly': True,
                    'message': f'已切换到数据库: {use_database}'
                }

            # 分析SQL类型
            sql_analysis = cls.analyze_sql_type(sql)

            # 获取数据库引擎
            engine = cls.get_engine(db_instance_to_use)

            # 执行查询
            with engine.connect() as conn:
                # 对于SELECT查询，添加LIMIT限制
                if sql_analysis['is_readonly'] and sql_analysis['sql_type'] == 'SELECT':
                    if 'LIMIT' not in sql.upper():
                        sql = f"{sql.rstrip(';')} LIMIT {limit}"

                # 对于非只读操作，使用事务
                if not sql_analysis['is_readonly']:
                    trans = conn.begin()
                    try:
                        result = conn.execute(text(sql))
                        trans.commit()
                    except Exception as e:
                        trans.rollback()
                        raise e
                else:
                    result = conn.execute(text(sql))

                # 处理查询结果
                if result.returns_rows:
                    # SELECT查询
                    columns = list(result.keys())
                    rows = []
                    row_count = 0

                    for row in result:
                        if row_count >= limit:
                            break
                        rows.append([str(value) if value is not None else '' for value in row])
                        row_count += 1

                    execution_time = time.time() - start_time

                    return {
                        'success': True,
                        'columns': columns,
                        'rows': rows,
                        'row_count': row_count,
                        'execution_time': round(execution_time, 3),
                        'sql_type': sql_analysis['sql_type'],
                        'is_readonly': sql_analysis['is_readonly'],
                        'message': f'查询执行成功，返回 {row_count} 行数据'
                    }
                else:
                    # DML/DDL操作
                    affected_rows = result.rowcount if hasattr(result, 'rowcount') else 0
                    execution_time = time.time() - start_time

                    return {
                        'success': True,
                        'columns': [],
                        'rows': [],
                        'row_count': 0,
                        'affected_rows': affected_rows,
                        'execution_time': round(execution_time, 3),
                        'sql_type': sql_analysis['sql_type'],
                        'is_readonly': sql_analysis['is_readonly'],
                        'message': f'{sql_analysis["sql_type"]} 执行成功，影响 {affected_rows} 行'
                    }

        except (OperationalError, ProgrammingError) as e:
            logger.error(f"SQL执行错误: {str(e)}")
            raise SQLExecutionError(f"SQL执行错误: {str(e)}")
        except SQLAlchemyError as e:
            logger.error(f"数据库操作错误: {str(e)}")
            raise SQLExecutionError(f"数据库操作错误: {str(e)}")
        except Exception as e:
            logger.error(f"未知错误: {str(e)}")
            raise SQLExecutionError(f"执行失败: {str(e)}")

    @classmethod
    def _execute_multiple_queries(cls, db_instance, statements: List[str], user, limit: int, start_time: float) -> Dict[str, Any]:
        """执行多条SQL查询"""
        results = []
        total_execution_time = 0
        successful_count = 0
        failed_count = 0

        # 检查是否有USE语句，如果有，应用到后续查询
        current_db_instance = db_instance

        for i, sql_content in enumerate(statements):
            sql = sql_content.strip()
            if not sql:
                continue

            try:
                # 为每条语句记录开始时间
                stmt_start_time = time.time()

                # 检查是否是USE语句
                sql_upper = sql.upper()
                if sql_upper.startswith('USE '):
                    # 处理USE语句
                    use_match = sql[4:].strip()  # 移除'USE '
                    use_database = use_match.rstrip(';').strip('`').strip("'").strip('"')
                    logger.info(f"检测到USE语句，切换到数据库: {use_database}")

                    # 创建一个临时的数据库实例副本，修改数据库名
                    from copy import copy
                    current_db_instance = copy(db_instance)
                    current_db_instance.database_name = use_database

                    stmt_execution_time = time.time() - stmt_start_time
                    total_execution_time += stmt_execution_time

                    results.append({
                        'success': True,
                        'columns': [],
                        'rows': [],
                        'row_count': 0,
                        'affected_rows': 0,
                        'execution_time': round(stmt_execution_time, 3),
                        'sql_type': 'USE',
                        'is_readonly': True,
                        'message': f'已切换到数据库: {use_database}',
                        'sql_content': sql
                    })
                    successful_count += 1
                    continue

                # 执行普通SQL语句
                result = cls._execute_single_query(current_db_instance, sql, user, limit, stmt_start_time)
                result['sql_content'] = sql
                results.append(result)

                stmt_execution_time = result.get('execution_time', 0)
                total_execution_time += stmt_execution_time
                successful_count += 1

            except Exception as e:
                stmt_execution_time = time.time() - stmt_start_time
                total_execution_time += stmt_execution_time
                failed_count += 1

                error_message = str(e)
                logger.error(f"第{i+1}条SQL执行失败: {error_message}")

                results.append({
                    'success': False,
                    'error': error_message,
                    'execution_time': round(stmt_execution_time, 3),
                    'sql_content': sql
                })

        # 返回多查询结果
        return {
            'success': True,
            'multiple_results': True,
            'results': results,
            'total_statements': len(statements),
            'successful_count': successful_count,
            'failed_count': failed_count,
            'total_execution_time': round(total_execution_time, 3),
            'message': f'批量执行完成：{successful_count}条成功，{failed_count}条失败'
        }
    
    @classmethod
    def get_database_schema(cls, db_instance) -> Dict[str, Any]:
        """获取数据库结构信息"""
        try:
            engine = cls.get_engine(db_instance)
            inspector = inspect(engine)
            
            # 获取所有表名
            tables = inspector.get_table_names()
            
            schema_info = {
                'tables': [],
                'total_tables': len(tables)
            }
            
            # 获取每个表的详细信息
            for table_name in tables[:50]:  # 限制返回前50个表
                try:
                    columns = inspector.get_columns(table_name)
                    table_info = {
                        'name': table_name,
                        'columns': [
                            {
                                'name': col['name'],
                                'type': str(col['type']),
                                'nullable': col.get('nullable', True),
                                'primary_key': col.get('primary_key', False)
                            }
                            for col in columns
                        ]
                    }
                    schema_info['tables'].append(table_info)
                except Exception as e:
                    logger.warning(f"获取表 {table_name} 信息失败: {str(e)}")
                    continue
            
            return schema_info
            
        except Exception as e:
            logger.error(f"获取数据库结构失败: {str(e)}")
            raise DatabaseConnectionError(f"获取数据库结构失败: {str(e)}")

    @classmethod
    def get_database_list(cls, db_instance, include_system_dbs: bool = False) -> Dict[str, Any]:
        """获取数据库实例下的数据库列表

        Args:
            db_instance: 数据库实例
            include_system_dbs: 是否包含系统数据库，默认False
        """
        try:
            engine = cls.get_engine(db_instance)

            # 根据数据库类型获取数据库列表
            if db_instance.database_type.lower() == 'mysql':
                sql = "SHOW DATABASES"
            elif db_instance.database_type.lower() == 'postgresql':
                sql = "SELECT datname as database_name FROM pg_database WHERE datistemplate = false"
            elif db_instance.database_type.lower() == 'sqlite':
                # SQLite 是文件数据库，没有多数据库概念
                return {
                    'databases': [{'name': 'main', 'type': 'database'}],
                    'total': 1
                }
            else:
                sql = "SHOW DATABASES"

            with engine.connect() as conn:
                result = conn.execute(text(sql))
                rows = result.fetchall()

                databases = []
                for row in rows:
                    db_name = row[0]

                    # 根据参数决定是否过滤系统数据库
                    if include_system_dbs:
                        # 显示所有数据库
                        databases.append({'name': db_name, 'type': 'database'})
                    else:
                        # 过滤系统数据库
                        if db_instance.database_type.lower() == 'mysql':
                            if db_name not in ['information_schema', 'performance_schema', 'mysql', 'sys']:
                                databases.append({'name': db_name, 'type': 'database'})
                        elif db_instance.database_type.lower() == 'postgresql':
                            if db_name not in ['template0', 'template1', 'postgres']:
                                databases.append({'name': db_name, 'type': 'database'})
                        else:
                            databases.append({'name': db_name, 'type': 'database'})

                return {
                    'databases': databases,
                    'total': len(databases)
                }

        except Exception as e:
            logger.error(f"获取数据库列表失败: {str(e)}")
            raise DatabaseConnectionError(f"获取数据库列表失败: {str(e)}")

    @classmethod
    def get_table_list(cls, db_instance, database_name) -> Dict[str, Any]:
        """获取指定数据库的表列表"""
        logger.info(f"get_table_list开始: db_instance={db_instance.name}, database_name={database_name}")

        try:
            # 验证参数
            if not database_name:
                raise DatabaseConnectionError("数据库名称不能为空")

            # 创建数据库实例副本，设置数据库名
            from copy import copy
            db_instance_copy = copy(db_instance)
            db_instance_copy.database_name = database_name
            logger.info(f"创建数据库实例副本，数据库名: {database_name}")

            # 获取连接字符串用于调试
            connection_string = cls.get_connection_string(db_instance_copy)
            logger.info(f"连接字符串: {connection_string}")

            engine = cls.get_engine(db_instance_copy)
            logger.info(f"成功获取数据库引擎")

            inspector = inspect(engine)
            logger.info(f"成功创建数据库检查器")

            # 获取所有表名
            tables = inspector.get_table_names()
            logger.info(f"成功获取到 {len(tables)} 个表")

            table_list = []
            for table_name in tables:
                table_list.append({
                    'name': table_name,
                    'type': 'table'
                })

            return {
                'tables': table_list,
                'total': len(table_list)
            }

        except Exception as e:
            import traceback
            logger.error(f"获取表列表失败: {str(e)}")
            logger.error(f"完整错误堆栈: {traceback.format_exc()}")
            raise DatabaseConnectionError(f"获取表列表失败: {str(e)}")

    @classmethod
    def get_table_structure(cls, db_instance, database_name, table_name) -> Dict[str, Any]:
        """获取表结构"""
        try:
            # 验证参数
            if not database_name:
                raise DatabaseConnectionError("数据库名称不能为空")
            if not table_name:
                raise DatabaseConnectionError("表名称不能为空")

            # 创建数据库实例副本，设置数据库名
            from copy import copy
            db_instance_copy = copy(db_instance)
            db_instance_copy.database_name = database_name
            logger.info(f"尝试获取表结构: {database_name}.{table_name}")

            engine = cls.get_engine(db_instance_copy)
            inspector = inspect(engine)

            # 获取表的列信息
            columns = inspector.get_columns(table_name)

            # 获取主键信息
            primary_keys = inspector.get_pk_constraint(table_name)

            # 获取索引信息
            indexes = inspector.get_indexes(table_name)

            # 获取外键信息
            foreign_keys = inspector.get_foreign_keys(table_name)

            table_structure = {
                'table_name': table_name,
                'columns': [
                    {
                        'name': col['name'],
                        'type': str(col['type']),
                        'nullable': col.get('nullable', True),
                        'default': col.get('default'),
                        'primary_key': col['name'] in primary_keys.get('constrained_columns', [])
                    }
                    for col in columns
                ],
                'primary_keys': primary_keys.get('constrained_columns', []),
                'indexes': [
                    {
                        'name': idx['name'],
                        'columns': idx['column_names'],
                        'unique': idx.get('unique', False)
                    }
                    for idx in indexes
                ],
                'foreign_keys': [
                    {
                        'name': fk['name'],
                        'constrained_columns': fk['constrained_columns'],
                        'referred_table': fk['referred_table'],
                        'referred_columns': fk['referred_columns']
                    }
                    for fk in foreign_keys
                ]
            }

            logger.info(f"成功获取表结构: {table_name}, 列数: {len(columns)}")
            return table_structure

        except Exception as e:
            logger.error(f"获取表结构失败: {str(e)}")
            raise DatabaseConnectionError(f"获取表结构失败: {str(e)}")

    @classmethod
    def close_all_connections(cls):
        """关闭所有数据库连接"""
        for engine in cls._engines.values():
            try:
                engine.dispose()
            except Exception as e:
                logger.error(f"关闭数据库连接失败: {str(e)}")
        cls._engines.clear()
